/**
* Leaflet地图MCP服务器
* 用于处理地图相关的MCP工具调用，并通过WebSocket发送消息给地图客户端
*/
import express, { Request, Response } from 'express';
import cors from 'cors';
import path from 'path';
import { fileURLToPath } from 'url';
import { WebSocketServer, WebSocket } from 'ws';
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { z } from "zod";
import * as turf from '@turf/turf';

// 路径兼容 CommonJS
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

// 初始化 Express 应用
const app = express();
app.use(cors());
app.use(express.static(path.join(__dirname)));

// 提供静态访问路径（可选）
app.get('/src/:filename', (req, res) => {
  res.sendFile(path.join(__dirname, 'src', req.params.filename));
});

// 启动 HTTP + WebSocket 服务
const server = app.listen(3000, () => {
  console.log("HTTP/WebSocket 服务器运行在 http://localhost:3000");
});

// 创建 WebSocket 服务端
const wss = new WebSocketServer({ server });
const clients = new Set<WebSocket>();

wss.on('connection', (ws) => {
  console.log('新的 WebSocket 客户端已连接');
  clients.add(ws);
  ws.on('close', () => clients.delete(ws));
});

function broadcastToClients(data: any) {
  const msg = JSON.stringify(data);
  for (const client of clients) {
    if (client.readyState === WebSocket.OPEN) {
      client.send(msg);
    }
  }
}

/**
 * 创建MCP服务器
 * 用于处理来自大模型的工具调用请求
 */
const mcpServer = new McpServer({
    name: "LeafletMapServer",
    version: "1.0.0"
});

/**
 * 定义坐标类型
 * 使用zod验证坐标是否为包含两个数字的数组
 */
const coordinateSchema = z.array(z.number()).length(2);

// 添加热力图工具
// 修改热力图工具以支持多个 geometry
mcpServer.tool("generateHeatmap", {
  geoms: z.array(z.object({
    type: z.enum(["Point", "LineString", "Polygon"]).describe("GeoJSON 几何类型：Point 点，LineString 线，Polygon 面"),
    coordinates: z.any().describe("GeoJSON 坐标结构。Point 为 [lng, lat]，LineString 为 [[lng, lat], ...]，Polygon 为 [[[lng, lat], ...]]"),
    intensity: z.number().optional().describe("每个图元的热力强度，可选。未指定时使用默认 intensity")
  })).describe("热力图图元数组，每个元素为 GeoJSON Geometry 对象，包含类型、坐标和可选的强度值"),
  
  intensity: z.number()
    .default(1.0)
    .describe("统一默认热力强度。用于未在 geoms 中单独指定 intensity 的情况"),
  
  heatmapOptions: z.object({
    radius: z.number().optional().describe("热力点扩散半径（像素），默认约 25"),
    blur: z.number().optional().describe("模糊半径（像素），默认约 15"),
    minOpacity: z.number().optional().describe("热力图最低透明度，默认 0"),
    maxZoom: z.number().optional().describe("热力图可见的最大缩放级别"),
    gradient: z.record(z.string()).optional().describe("热力图颜色渐变映射，key 为 0~1 的数值字符串，value 为颜色名或 hex 值")
  }).optional().describe("热力图图层样式配置项，支持 Leaflet heatLayer 可用参数")
}, async ({ geoms, intensity: defaultIntensity, heatmapOptions }) => {
  const heatPoints: [number, number, number][] = [];

  for (const geom of geoms) {
    if (!geom) continue;

    // 兜底逻辑：每个 geometry 的 intensity 优先，没有则用统一默认值
    const intensity = geom.intensity ?? defaultIntensity;

    if (geom.type === 'Point') {
      const [lng, lat] = geom.coordinates;
      heatPoints.push([lat, lng, intensity]);
    }
    else if (geom.type === 'LineString') {
      for (const [lng, lat] of geom.coordinates) {
        heatPoints.push([lat, lng, intensity]);
      }
    }
    else if (geom.type === 'Polygon') {
      const polygon = turf.polygon(geom.coordinates);
      const sampled = turf.pointGrid(turf.bbox(polygon), 0.01, { mask: polygon });
      for (const pt of sampled.features) {
        const [lng, lat] = pt.geometry.coordinates;
        heatPoints.push([lat, lng, intensity]);
      }
    }
  }

  // 向前端广播热力图数据和图层样式配置
  broadcastToClients({
    success: true,
    command: "generateHeatmap",
    params: {
      heatPoints,
      heatmapOptions
    }
  });

  return {
    content: [{
      type: "text",
      text: JSON.stringify({ success: true, command: "generateHeatmap" })
    }]
  };
});



mcpServer.tool("init", {
  center: coordinateSchema.describe("中心坐标点，格式为 [lat, lng]，即 [纬度, 经度]"), 
  zoom: z.number().default(13).describe("地图缩放级别,默认为 13"),
}, async ({ center, zoom }) => {

  // 广播给客户端
  broadcastToClients({
    success: true,
    command: "init",
    params: { center, zoom }
  });

  return {
    content: [{
      type: "text",
      text: JSON.stringify({
        success: true,
        command: "init",
        params: { center, zoom }
      })
    }]
  };
});

mcpServer.tool("addMarker", {
  coordinate: coordinateSchema.describe("坐标点，格式为 [lat, lng]，即 [纬度, 经度]"),
  color: z.string().default("#3388ff").describe("线条颜色，支持任意合法的 CSS 颜色格式，如 '#ff0000' 或 'blue'"),
  size: z.number().default(5).describe("标记大小（像素）"),
}, async ({ coordinate, color, size }) => {
  broadcastToClients({ success: true, command: "addMarker", params: { coordinate, color, size } });
  return { content: [{ type: "text", text: JSON.stringify({ success: true, command: "addMarker", params: { coordinate, color, size } }) }] };
});

mcpServer.tool("clearLayers", {
  layerType: z.string().default("all").describe("图层类型，可选 'markers'、'lines'、'polygons' 或 'all'，默认 'all'，清除所有图层")
}, async ({ layerType }) => {
  broadcastToClients({ success: true, command: "clearLayers", params: { layerType } });
  return { content: [{ type: "text", text: JSON.stringify({ success: true, command: "clearLayers", params: { layerType } }) }] };
});

mcpServer.tool("drawLine", {
  coordinates: z.array(coordinateSchema).min(2).describe("坐标点数组，格式为 [[lat, lng], [lat, lng], ...]，即 [[纬度, 经度], ...]"),
  color: z.string().default("#3388ff").describe("线条颜色，支持任意合法的 CSS 颜色格式，如 '#ff0000' 或 'blue'"),
  weight: z.number().default(3).describe("线的宽度（像素）"),
}, async ({ coordinates, color, weight }) => {
  broadcastToClients({ success: true, command: "drawLine", params: { coordinates, color, weight } });
  return { content: [{ type: "text", text: JSON.stringify({ success: true, command: "drawLine", params: { coordinates, color, weight } }) }] };
});

mcpServer.tool("drawPolygon", {
  coordinates: z.array(coordinateSchema).min(3).describe("坐标点数组，格式为 [[lat, lng], [lat, lng], ...]，即 [[纬度, 经度], ...]"),
  color: z.string().default("#3388ff").describe("线条颜色，支持任意合法的 CSS 颜色格式，如 '#ff0000' 或 'blue'"),
  fillColor: z.string().default("#3388ff").describe("多边形的填充色，默认与边框色相同"),
  weight: z.number().default(3).describe("线的宽度（像素）"),
}, async ({ coordinates, color, fillColor, weight }) => {
  broadcastToClients({ success: true, command: "drawPolygon", params: { coordinates, color, fillColor, weight } });
  return { content: [{ type: "text", text: JSON.stringify({ success: true, command: "drawPolygon", params: { coordinates, color, fillColor, weight } }) }] };
});

// 坐标转换工具
mcpServer.tool("coordinateTransform", {
  coordinates: z.array(coordinateSchema).describe("需要转换的坐标数组，格式为[[lng, lat], [lng, lat], ...]"),
  from: z.enum(["WGS84", "GCJ02", "BD09"]).describe("原始坐标系"),
  to: z.enum(["WGS84", "GCJ02", "BD09"]).describe("目标坐标系")
}, async ({ coordinates, from, to }) => {
  if (from === to) {
    return {
      content: [{
        type: "text",
        text: JSON.stringify({
          success: true,
          command: "coordinateTransform",
          result: coordinates
        })
      }]
    };
  }

  broadcastToClients({
    success: true,
    command: "coordinateTransform",
    params: { coordinates, from, to }
  });

  return {
    content: [{
      type: "text",
      text: JSON.stringify({
        success: true,
        command: "coordinateTransform",
        message: "坐标转换请求已发送到客户端处理"
      })
    }]
  };
});

mcpServer.tool("boundaryLayer", {
  province: z.string().describe("区域名称，例如 '天河区'、'东莞市'、'广东省' 等，支持模糊匹配")
}, async ({ province }) => {
  try {
    const provinceGeoJSON = await getProvinceBoundary(province);
    broadcastToClients({
      success: true,
      command: "drawBoundary",
      params: { province, provinceGeoJSON }
    });
    return {
      content: [{
        type: "text",
        text: JSON.stringify({
          success: true,
          command: "drawBoundary",
          params: { province }
        })
      }]
    };
  }
  catch (error) {
    const errorMessage = error instanceof Error ? error.message : String(error);
    return {
      content: [{
        type: "text",
        text: JSON.stringify({
          success: false,
          error: `获取${province}边界数据失败: ${errorMessage}`
        })
      }]
    };
  }
});

mcpServer.tool("zoomMap", {
  center: coordinateSchema.describe("地图中心点，格式为 [lat, lng]，即 [纬度, 经度]"),
  zoom: z.number().describe("地图缩放级别")
}, async ({ center, zoom }) => {
  broadcastToClients({ success: true, command: "zoomMap", params: { center, zoom } });
  return {
    content: [{
      type: "text",
      text: JSON.stringify({ success: true, command: "zoomMap", params: { center, zoom } })
    }]
  };
});




// 保留getProvinceBoundary函数和SSE服务端代码...
async function getProvinceBoundary(areaName: string) {
  const apiUrl = `https://restapi.amap.com/v3/config/district?keywords=${encodeURIComponent(areaName)}&subdistrict=1&extensions=all&key=15f921f471f73ffedb54b93d1cdb9773`;

  try {
    const response = await fetch(apiUrl);
    if (!response.ok) {
      throw new Error(`请求失败: ${response.status} ${response.statusText}`);
    }

    const data = await response.json();
    if (!data.districts || data.districts.length === 0) {
      throw new Error(`未找到 ${areaName} 的边界数据`);
    }

    const findMatchingDistrict = (districts: any[]): any => {
      for (const district of districts) {
        if (district.name.includes(areaName)) {
          return district;
        }
        if (district.districts) {
          const found = findMatchingDistrict(district.districts);
          if (found) return found;
        }
      }
      return null;
    };

    const matchedDistrict = findMatchingDistrict(data.districts);
    if (!matchedDistrict || !matchedDistrict.polyline) {
      throw new Error(`未找到 ${areaName} 的边界数据`);
    }

    const coordinates = matchedDistrict.polyline.split('|').map((polygon: string) => {
      const rings = polygon.split(';').map((point: string) => {
        const [lng, lat] = point.split(',').map(Number);
        return [lng, lat];
      });
      if (rings.length > 0 && !(rings[0][0] === rings[rings.length - 1][0] && rings[0][1] === rings[rings.length - 1][1])) {
        rings.push(rings[0]);
      }
      return [rings];
    });

    const feature = {
      type: "Feature",
      properties: { name: areaName },
      geometry: {
        type: "MultiPolygon",
        coordinates: coordinates
      }
    };

    const geoJSON = {
      type: "FeatureCollection",
      features: [feature]
    };

    try {
      const geoJSONString = JSON.stringify(geoJSON);
      JSON.parse(geoJSONString);
    } catch (error) {
      console.error('生成的GeoJSON格式错误:', error);
      throw new Error('生成的GeoJSON格式不符合规范');
    }

    return geoJSON;
  } catch (error) {
    throw new Error(`获取 ${areaName} 边界数据失败: ${error instanceof Error ? error.message : String(error)}`);
  }
}

// 启动SSE服务端
const transports: { [sessionId: string]: SSEServerTransport } = {};

app.get("/sse", async (_: Request, res: Response) => {
  const transport = new SSEServerTransport('/messages', res);
  transports[transport.sessionId] = transport;
  res.on("close", () => {
    delete transports[transport.sessionId];
  });
  await mcpServer.connect(transport);
});

app.post("/messages", async (req: Request, res: Response) => {
  const sessionId = req.query.sessionId as string;
  const transport = transports[sessionId];
  if (transport) {
    await transport.handlePostMessage(req, res);
  } else {
    res.status(400).send('No transport found for sessionId');
  }
});

app.listen(8001, () => {
  console.log("Server is running on http://localhost:8001");
});
